package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{Word2VecEncoder, Word2VecEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.sql.functions.udf
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2019/3/26 14:32
  * @Description
  * @Reviewer
  */
class Word2VecEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "sentence",
      |	"maxIter": 4,
      |	"outputCol": "vec",
      |	"preserveCols": "sentence"
      |}
    """.stripMargin

  val w2v = new Word2VecEncoder
  val param: Word2VecEncoderParam = w2v.parseParam(new ToolClass().encrypt(paramStr))

  test("w2v parameter") {
    assert("sentence".equals(param.inputCol) && "vec".equals(param.outputCol) &&
      "sentence".equals(param.preserveCols) && param.maxIter == 4 && param.minCount == 5 && param.vectorSize == 100 &&
      param.maxSentenceLength == 1000 && param.numPartitions == 1 && param.seed == 123456L && param.stepSize == 0.025 &&
      param.windowSize == 5)
  }


  test("w2v transformation") {
    var data = TestingDataGenerator.pushkinPoem
    val spliter = udf{ str: String => {
      str.split(" ")
    }}
    data = data.withColumn(param.inputCol, spliter(data.col("sentence")))
    val res = w2v.train(param, data).transform(data)
    // 如果输出列是 Vector 类型，说明组件运行正常
    assert(res.select(param.outputCol).rdd.map(_.get(0)).first.isInstanceOf[org.apache.spark.ml.linalg.Vector])
  }

}
